#!/usr/bin/env python

import rospy
from unitree_legged_msgs.msg import LowState
from config import Config
import numpy as np


# 指定文件名和模式 'a' 表示追加模式

def update_linear_velocity(current_velocity,current_acceleration, time_interval):

   
    # Update velocity using Euler method
    updated_velocity = current_velocity + current_acceleration * time_interval
 
    return updated_velocity


def write2txt(data):
    filename = "data.txt" # 要写入的数据 
    # data = "这是一行新的数据." # 打开文件并写入数据 
    with open(filename, 'a', encoding='utf-8') as file: # 写入数据，自动换行 
        file.write(data+"\n") # 如果需要在数据之间插入换行符，可以这样写： 

class LowStateListener:
    def __init__(self):
        self.motor_state={}
        self.joint_positions = {}
        self.joint_velocities = {}
        self.oritation={}
        self.data=[0]*13
        self.current_velocity = np.array([0.0, 0.0, 0.0])
        self.interval=1/Config.frequency

        rospy.Subscriber("/low_state_info", LowState, self.callback)

    def callback(self, msg ):
        # imuData_.ori_[0] = msg.imu.quaternion[1]
        # imuData_.ori_[1] = msg.imu.quaternion[2]
        # imuData_.ori_[2] = msg.imu.quaternion[3]
        # imuData_.ori_[3] = msg.imu.quaternion[0]
        # imuData_.angularVel_[0] = msg.imu.gyroscope[0]
        # imuData_.angularVel_[1] = msg.imu.gyroscope[1]
        # imuData_.angularVel_[2] = msg.imu.gyroscope[2]
        # imuData_.linearAcc_[0] = msg.imu.accelerometer[0]
        # imuData_.linearAcc_[1] = msg.imu.accelerometer[1]
        # imuData_.linearAcc_[2] = msg.imu.accelerometer[2]
        accleration=np.array([msg.imu.accelerometer[0],msg.imu.accelerometer[1],msg.imu.accelerometer[2]])
        accleration[2]=accleration[2]-9.6

        self.current_velocity=update_linear_velocity(self.current_velocity,accleration,self.interval)
        # self.data=[0, 0, 0,  msg.imu.quaternion[1] ,msg.imu.quaternion[2], msg.imu.quaternion[3], msg.imu.quaternion[0],
        #     msg.velocity.x, msg.velocity.y, msg.velocity.z,   msg.imu.gyroscope[0], msg.imu.gyroscope[1], msg.imu.gyroscope[2]  ]
        self.data=[0, 0, 0,  msg.imu.quaternion[1] ,msg.imu.quaternion[2], msg.imu.quaternion[3], msg.imu.quaternion[0],
           self.current_velocity[0], self.current_velocity[1], self.current_velocity[2],   msg.imu.gyroscope[0], msg.imu.gyroscope[1], msg.imu.gyroscope[2]  ]
        
        # write2txt("x: {} y:{}  z:{}".format(self.current_velocity[0], self.current_velocity[1], self.current_velocity[2]))
        for joint_name, index in  Config.joint2MotorIndex.items():
            self.joint_positions[joint_name] = msg.motorState[index].q
            self.joint_velocities[joint_name] = msg.motorState[index].dq

    def get_joint_states(self):         
        return self.joint_positions, self.joint_velocities

    def get_model_states(self):         
        # position = Point4D(0,0,0,)
        # orientation = data.pose[i].orientation
        # linear_velocity = data.twist[i].linear
        # angular_velocity = data.twist[i].angular
        # # rospy.loginfo("  Position: x=%f, y=%f, z=%f", position.x, position.y, position.z)
        # # rospy.loginfo("  Orientation: x=%f, y=%f, z=%f, w=%f", orientation.x, orientation.y, orientation.z, orientation.w)
        # # rospy.loginfo("  Linear Velocity: x=%f, y=%f, z=%f", linear_velocity.x, linear_velocity.y, linear_velocity.z)
        # # rospy.loginfo("  Angular Velocity: x=%f, y=%f, z=%f", angular_velocity.x, angular_velocity.y, angular_velocity.z)
        # self.data=[position.x, position.y, position.z, orientation.x, orientation.y, orientation.z, orientation.w,
        #             linear_velocity.x, linear_velocity.y, linear_velocity.z,  angular_velocity.x, angular_velocity.y, angular_velocity.z  ]
        return self.data


if __name__ == "__main__":
    rospy.init_node('joint_state_listener_node', anonymous=True)
    listener = LowStateListener()
    
    rate = rospy.Rate(60)  # 60xjz Hz
    while not rospy.is_shutdown():
        positions, velocities = listener.get_joint_states()
        rospy.loginfo("Joint Positions: {}".format(positions))
        rospy.loginfo("Joint Velocities: {}".format(velocities))
        rate.sleep()
